from elasticsearch import Elasticsearch

import re

import stixconfig
from stix.core import STIXPackage, ExploitTarget
from stix.report import Report
from stix.indicator import Indicator, Observable
from stix.exploit_target.vulnerability import Vulnerability, CVSSVector
from stix.common.datetimewithprecision import DateTimeWithPrecision
from stix.common import InformationSource
from stix.common.information_source import ContributingSources
from mixbox.namespaces import Namespace
from mixbox.idgen import set_id_namespace
from cybox.objects.file_object import File, FilePath
from cybox.objects.system_object import OS, System
from cybox.objects.code_object import Code
from cybox.core.observable import ObservableComposition, Observables
from cybox.common import Hash
from cybox.common.platform_specification import PlatformSpecification, PlatformIdentifier
from maec.bundle import Bundle, AVClassification

set_id_namespace(Namespace(stixconfig.namespace, stixconfig.domain))



def format_query_body(field, data):
    q_body = {
        'size': stixconfig.limit_count,
        'timeout': stixconfig.timeout,
        'query': {
            'filtered':
                {
                    'query': {'match_all': {}},
                    'filter': {
                        'bool':
                            {
                                'must': [{'term': {field: data}}]
                            }
                    }
                }
        }
    }
    return q_body

class stixhandler:
    def __init__(self, body=None, index=None, doc_type=None):
        self.body = body
        self.index=index
        self.doc_type=doc_type

        self.es = Elasticsearch(stixconfig.es_url)

        # search cache
        cache_body ={
            'query': {
                'match': {
                    'body': body,
                    'index': index,
                    'doc_type': doc_type
                }
            }
        }
        try:
            cache_res = self.es.search(index='xmlcache', doc_type='package',body=cache_body)
        except:
            cache_res={'hits':{'total':0}}

        if cache_res['hits']['total'] > 0:
            self.xml = cache_res['hits']['hits'][-1]['_source']['result']
        else:
            self.stix_package = STIXPackage()
            self.searchines(index=index, doc_type=doc_type, body=body)
            self.xml = self.stix_package.to_xml()
            # save cache
            resbody = {
                'body': body,
                'index': index,
                'doc_type':doc_type,
                'result':self.xml
            }
            self.es.index(index='xmlcache', doc_type='package', body=resbody)


    # recurisive
    def searchines(self, body=None, index=None, doc_type=None, relate_parent=None):
        resdict = self.es.search(body=body, index=index, doc_type=doc_type)
        for doc_dict in resdict['hits']['hits']:
            if doc_dict['_index'] == 'cve':
                if doc_dict['_type'] == 'CVEs':
                    self.format_CVEs(doc_dict, relate_parent=relate_parent)
                if doc_dict['_type'] == 'PPSs':
                    self.format_PPSs(doc_dict)
                if doc_dict['_type'] == 'EDBs':
                    self.format_EDBs(doc_dict)
            if doc_dict['_index'] == 'ttp':
                if doc_dict['_type'] == 'DNA':
                    self.format_DNA(doc_dict)
                if doc_dict['_type'] == 'AVs':
                    self.format_AVs(doc_dict)

    def format_CVEs(self, doc_dict, relate_parent=None):
        report = Report()

        exp_target = ExploitTarget()
        vuln = Vulnerability()
        vuln.description = doc_dict['_source']['descript']
        # cve
        vuln.cve_id = doc_dict['_source']['cvename']
        # cvss
        cvss_vector = CVSSVector()
        cvss_vector.base_vector = doc_dict['_source']['CVSS_vector'].lstrip('(').rstrip(')')
        cvss_vector.base_score = doc_dict['_source']['CVSS_base_score']
        cvss_vector.overall_score = doc_dict['_source']['CVSS_score']
        vuln.cvss_score = cvss_vector
        # date
        vuln.published_datetime = DateTimeWithPrecision(value=doc_dict['_source']['published'], precision='day')
        ###
        exp_target.add_vulnerability(vuln)
        report.add_exploit_target(exp_target)

        if relate_parent:
            relate_parent.add_related_report(report)
        else:
            self.stix_package.add_report(report)

    def format_PPSs(self, doc_dict):
        report = Report()
        indicator = Indicator()
        observables = Observables()

        contribsources = ContributingSources()
        contribsources.append(doc_dict['_source']['source'])
        infosource = InformationSource(identity=doc_dict['_source']['credit'], contributing_sources=contribsources)
        indicator.producer = infosource
        indicator.producer.time = doc_dict['_source']['published']

        obsf = File()
        obsf.title = doc_dict['_source']['title']
        obsf.file_name = doc_dict['_source']['title']
        obsf.file_path = FilePath(doc_dict['_source']['downloadlink'])
        obsf.full_path = FilePath(doc_dict['_source']['filelink'])
        obs = Observable(item=obsf, id_=doc_dict['_source']['fid'])
        obs.keywords.from_list(doc_dict['_source']['tags'].split(','))
        obs.description = doc_dict['_source']['description']
        observables.add(obs)

        obss = System()
        obso = OS()
        platform_spec = PlatformSpecification()
        platform_spec.description = doc_dict['_source']['systems']
        obso.platform = platform_spec
        obss.os = obso
        obs = Observable(item=obss)
        observables.add(obs)

        obs_comp = ObservableComposition(operator='AND', observables=observables)
        indicator.add_observable(obs_comp)
        report.id_ = doc_dict['_source']['id']
        report.add_indicator(indicator)

        cve_pattern = re.compile(r'(CVE-\d{4}-\d{4,7})')
        cve_ids = cve_pattern.findall(doc_dict['_source']['description'])
        for cve_id in cve_ids:
            self.searchines(body=format_query_body(field='cvename', data=cve_id), relate_parent=report)

        self.stix_package.add_report(report)

    def format_EDBs(self, doc_dict):
        report = Report()
        indicator = Indicator()
        observables = Observables()

        contribsources = ContributingSources()
        contribsources.append(doc_dict['_source']['source'])
        infosource = InformationSource(identity=doc_dict['_source']['credit'], contributing_sources=contribsources)
        indicator.producer = infosource
        indicator.producer.time = doc_dict['_source']['published']

        obsc = Code()
        obsc.code_language = doc_dict['_source']['filetype']
        observables.add(obsc)

        obsf = File()
        obsf.file_path = FilePath(doc_dict['_source']['downloadlink'])
        obsf.full_path = FilePath(doc_dict['_source']['filelink'])
        obs = Observable(item=obsf, id_=doc_dict['_source']['fid'])
        obs.keywords.from_list(doc_dict['_source']['tags'].split(','))
        obs.description = doc_dict['_source']['title']
        observables.add(obs)

        obs_comp = ObservableComposition(operator='AND', observables=observables)
        indicator.add_observable(obs_comp)
        report.id_ = doc_dict['_source']['id']
        report.add_indicator(indicator)

        if doc_dict['_source']['ereferences']:
            self.searchines(body=format_query_body(field='cvename', data=doc_dict['_source']['ereferences']), relate_parent=report)

        self.stix_package.add_report(report)

    def format_DNA(self, doc_dict, relate_parent=None):
        report = Report()
        indicator = Indicator()
        observables = Observables()

        indicator.description = doc_dict['_source']['vtresults']

        obsc = Code()
        obsc.type_ = doc_dict['_source']['dtype']
        obsc.description = doc_dict['_source']['vtresults']
        observables.add(obsc)

        obsf = File()
        md5hash = Hash(hash_value=doc_dict['_source']['md5'], type_='MD5')
        ssdeephash = Hash(hash_value=doc_dict['_source']['ssdeep'], type_='SSDEEP')
        obsf.add_hash(md5hash)
        obsf.add_hash(ssdeephash)
        obsf.size_in_bytes = doc_dict['_source']['fsize']
        obsf.created_time = doc_dict['_source']['create_at']
        observables.add(obsf)

        obs_comp = ObservableComposition(operator='AND', observables=observables)
        indicator.add_observable(obs_comp)
        report.id_ = doc_dict['_source']['id']
        report.add_indicator(indicator)

        if relate_parent:
            relate_parent.add_related_report(report)
        else:
            self.stix_package.add_report(report)

    def format_AVs(self, doc_dict):
        report = Report()

        maecbundle = Bundle(malware_instance_object=None)
        for av in doc_dict['_source']['detail']:
            avc = AVClassification()
            avc.name = av['scanEngine']
            avc.classification_name = av['report']
            maecbundle.add_av_classification(avc)

        report.add(maecbundle)
        self.stix_package.add_report(report)

    def exportxml(self):
        return self.xml
